#ifndef _BEZIER_H_
#define _BEZIER_H_

#include <Eigen/Eigen>
#include <vector>
#include "OsqpEigen/OsqpEigen.h"

using std::vector;

class Bezierpredict
{
private:
  Eigen::VectorXd PolyCoeff_;
  double poly_start_; //  relative start time of prediction on the Bezier curve
  double poly_end_;
  Eigen::MatrixXd M_;  // 映射矩阵
  Eigen::MatrixXd S_;  // smooth matrix
  std::vector<int> C_; // 组合数

  double time_history_{0.0}; // = poly_start_ if fit success
  double time_scale_{2.2};   // time total = time_scale_ * time_history_
  double time_total_{0.0};   // = poly_end_ if fit success

  int TRAJ_ORDER{5};           // 多项式阶数
  int SEGS_MAX{30};            // 历史轨迹段数
  int SEGS_MIN{7};             // 最小轨迹段数
  double SAMPLE_INTERVAL{0.1}; // 采样间隔
  double Lambda_ACC{0.05};     // 最优化系数

public:
  double t_predict_start_gl_;  // global start time of prediction
  double t_start_gl_;          // global start time
  double R2_;                  // adjust R2, goodness of fit
  double prop_start_;          // propagation coefficient at time = poly_start_
  Eigen::MatrixXd Evaluation_; // for calculate the variance of prediction
  Eigen::Vector3d Sigma2_;     // variance

  int factorial(int n)
  {
    int fact = 1;
    for (int i = n; i > 0; i--)
      fact *= i;
    return fact;
  }
  int combinatorial(int n, int k)
  {
    return factorial(n) / (factorial(k) * factorial(n - k));
  }

  Bezierpredict()
  {
    M_ = Eigen::MatrixXd::Zero(6, 6);
    M_ << 1, 0, 0, 0, 0, 0,
        -5, 5, 0, 0, 0, 0,
        10, -20, 10, 0, 0, 0,
        -10, 30, -30, 10, 0, 0,
        5, -20, 30, -20, 5, 0,
        -1, 5, -10, 10, -5, 1;
    for (int i = 0; i < TRAJ_ORDER + 1; i++) {
      C_.push_back(combinatorial(TRAJ_ORDER, i));
    }

    S_ = getS(TRAJ_ORDER + 1, 2);
  }

  ~Bezierpredict() {}

  void setSegInterval(int min, int max)
  {
    if (min > max)
      std::swap(min, max);
    if (min < TRAJ_ORDER + 2)
      std::cerr << "The min segs number should > TRAJ_ORDER+2" << std::endl;
    SEGS_MIN = min;
    SEGS_MAX = max;
  }

  void setSampleInterval(double interval) { SAMPLE_INTERVAL = interval; }
  void setTimeScale(double scale) { time_scale_ = scale; }

  bool TrackingGeneration(
      const double max_vel,
      const double max_acc,
      vector<Eigen::Vector4d> predict_list_complete);

  Eigen::MatrixXd getM(const int vars_number, double Time)
  {
    Eigen::MatrixXd M_k = Eigen::MatrixXd::Zero(vars_number, vars_number);
    Eigen::VectorXd t_pow = Eigen::VectorXd::Zero(vars_number);
    for (int i = 0; i < vars_number; i++)
    {
      t_pow(i) = pow(Time, i);
    }
    M_k = M_;
    for (int i = 0; i < vars_number; i++) {
      M_k.row(i) = M_k.row(i) / t_pow(i);
    }
    return M_k;
  }

  Eigen::MatrixXd getS(int vars_number, int derivative_order)
  {
    if (derivative_order >= vars_number) {
      std::cout << "[ERROR]:derivative should < vars_number" << std::endl;
      return Eigen::MatrixXd::Zero(vars_number, vars_number);
    }

    Eigen::Vector2d d(-1, 1);
    Eigen::MatrixXd D = Eigen::MatrixXd::Zero(vars_number - 1, vars_number);
    Eigen::MatrixXd S;
    for (int i = 0; i < D.rows(); i++)
    {
      D.block(i, i, 1, 2) = d.transpose();
    }

    S = D * (vars_number - 1);

    for (int i = 0; i < derivative_order - 1; i++)
    {
      S = D.block(0, 0, S.rows() - 1, S.rows()) * S * (vars_number - 2 - i);
    }

    return S.transpose() * S;
  }

  inline Eigen::MatrixXd getCt(double curr_time, Eigen::Vector3d curr_pos)
  {
    Eigen::MatrixXd Ct = Eigen::MatrixXd::Zero(1, 3 * (TRAJ_ORDER + 1));
    for (int i = 0; i < 3; i++) { // X Y Z
      for (int j = 0; j < TRAJ_ORDER + 1; j++) {
        Ct(0, j + i * (1 + TRAJ_ORDER)) = (-2) * curr_pos[i] * pow(curr_time, j);
      }
    }
    return Ct;
  }

  Eigen::VectorXd getT(double t)
  {
    Eigen::VectorXd T = Eigen::VectorXd::Zero(TRAJ_ORDER + 1);
    for (int i = 0; i < TRAJ_ORDER + 1; i++)
    {
      T(i) = pow(t, i);
    }
    return T;
  }

  inline Eigen::MatrixXd getTT(double curr_time)
  {
    Eigen::MatrixXd TT = Eigen::MatrixXd::Zero(TRAJ_ORDER + 1, TRAJ_ORDER + 1);
    for (int j = 0; j < TRAJ_ORDER + 1; j++) {
      for (int k = 0; k < TRAJ_ORDER + 1; k++) {
        TT(j, k) = pow(curr_time, j + k);
      }
    }
    return TT;
  };

  inline Eigen::Vector3d getPos(const double &t_now)
  {
    Eigen::Vector3d ret = Eigen::VectorXd::Zero(3);
    double T = poly_end_;

    for (int i = 0; i < 3; i++)
      for (int j = 0; j < TRAJ_ORDER + 1; j++)
        ret(i) += C_[j] * PolyCoeff_(i * (TRAJ_ORDER + 1) + j) *
                  pow(t_now / T, j) * pow((1 - t_now / T), (TRAJ_ORDER - j));
    return ret;
  };

  inline Eigen::Vector3d getVel(const double &t_now)
  {
    Eigen::Vector3d ret = Eigen::VectorXd::Zero(3);
    double T = poly_end_;
    for (int i = 0; i < 3; i++)
      for (int j = 0; j < TRAJ_ORDER; j++)
      {
        double c_n =
            TRAJ_ORDER *
            (PolyCoeff_(i * (TRAJ_ORDER + 1) + j + 1) -
             PolyCoeff_(i * (TRAJ_ORDER + 1) + j)) /
            T;
        ret(i) += combinatorial(TRAJ_ORDER - 1, j) * c_n * pow(t_now / T, j) *
                  pow((1 - t_now / T), (TRAJ_ORDER - 1 - j));
      }
    return ret;
  };

  inline Eigen::Vector3d getAcc(const double &t_now)
  {
    Eigen::Vector3d ret = Eigen::VectorXd::Zero(3);
    double T = poly_end_;
    for (int i = 0; i < 3; i++)
      for (int j = 0; j < TRAJ_ORDER - 1; j++)
      {
        double c_n =
            TRAJ_ORDER * (TRAJ_ORDER - 1) *
            (PolyCoeff_(i * (TRAJ_ORDER + 1) + j + 2) -
             2 * PolyCoeff_(i * (TRAJ_ORDER + 1) + j + 1) +
             PolyCoeff_(i * (TRAJ_ORDER + 1) + j)) /
            pow(T, 2);
        ret(i) += combinatorial(TRAJ_ORDER - 2, j) * c_n * pow(t_now / T, j) *
                  pow((1 - t_now / T), (TRAJ_ORDER - 2 - j));
      }
    return ret;
  };

  inline Eigen::Vector3d getJerk(const double &t_now)
  {
    Eigen::Vector3d ret = Eigen::VectorXd::Zero(3);
    double T = poly_end_;
    for (int i = 0; i < 3; i++)
      for (int j = 0; j < TRAJ_ORDER - 2; j++)
      {
        double c_n =
            TRAJ_ORDER * (TRAJ_ORDER - 1) * (TRAJ_ORDER - 2) *
            (PolyCoeff_(i * (TRAJ_ORDER + 1) + j + 3) -
             3 * PolyCoeff_(i * (TRAJ_ORDER + 1) + j + 2) +
             3 * PolyCoeff_(i * (TRAJ_ORDER + 1) + j + 1) -
             PolyCoeff_(i * (TRAJ_ORDER + 1) + j)) /
            pow(T, 3);
        ret(i) += combinatorial(TRAJ_ORDER - 3, j) * c_n * pow(t_now / T, j) *
                  pow((1 - t_now / T), (TRAJ_ORDER - 3 - j));
      }
    return ret;
  };

  Eigen::Vector3d getStartVel()
  {
    return getVel(poly_start_);
  };

  Eigen::Vector3d getStartAcc()
  {
    return getAcc(poly_start_);
  };

  inline vector<Eigen::Vector3d> getPredictPosList()
  {
    vector<Eigen::Vector3d> PosList;
    for (double t = poly_start_; t < poly_end_; t += SAMPLE_INTERVAL) {
      Eigen::Vector3d pos = getPos(t);
      PosList.push_back(pos);
    }

    return PosList;
  }

  inline vector<Eigen::Matrix<double, 6, 1>> getStateList()
  {
    vector<Eigen::Matrix<double, 6, 1>> predict_traj_state;
    // 预测时间到预测结束时间
    for (double t = poly_start_; t < poly_end_; t += SAMPLE_INTERVAL) {
      Eigen::Vector3d pos = getPos(t);
      Eigen::Vector3d vel = getVel(t);
      Eigen::Matrix<double, 6, 1> tmp_state;
      tmp_state.head(3) = pos;
      tmp_state.tail(3) = vel;
      predict_traj_state.push_back(tmp_state);
    }

    return predict_traj_state;
  }

  /**
   * @param t_now global time
   * @param R2 goodness of fit
   * @param propagation_coef propagation coefficient
   */
  void getUncertainty(double t_now, double &R2)
  {
    double t = t_now - t_start_gl_;
    Eigen::VectorXd T = getT(t);
    double propagation_coef = T.dot(Evaluation_ * T) / prop_start_;
    R2 = 1 - propagation_coef * (1 - R2_);
  }

  std::vector<Eigen::Vector3d> getControlPoints()
  {
    if (PolyCoeff_.rows() != 3 * (TRAJ_ORDER + 1)) {
      return std::vector<Eigen::Vector3d>();
    }

    std::vector<Eigen::Vector3d> ctrl_pts(TRAJ_ORDER + 1);
    for (size_t i = 0; i < TRAJ_ORDER + 1; i++)
    {
      ctrl_pts[i] = {PolyCoeff_(i),
                     PolyCoeff_(i + TRAJ_ORDER + 1),
                     PolyCoeff_(i + 2 * TRAJ_ORDER + 2)};
    }
    return ctrl_pts;
  }

  double getR2() { return R2_; }
};

#endif